import argparse
import os
import sys
import cv2
import torch
import numpy as np
from PIL import Image
import matplotlib.pylab as plt
sys.path.append(".")
from stereoflow.test import _load_model_and_criterion
from stereoflow.engine import tiled_pred
from stereoflow.datasets_stereo import img_to_tensor
from utils.config import *

use_gpu = torch.cuda.is_available() and torch.cuda.device_count()>0
device = torch.device('cuda:0' if use_gpu else 'cpu')
tile_overlap=0.7 # recommended value, higher value can be slightly better but slower

def get_args_parser():
    parser = argparse.ArgumentParser(description='Demo')

    parser.add_argument('--loadmodel', type=str, default='stereoflow_models/crocostereo.pth', help='load model')
    parser.add_argument('--dataset', type=str, default='kitti', help='dataset', choices=['kitti', 'usvinland', 'usvinland_seg'])
    parser.add_argument('--imgname', type=str, default='000000_10', help='image name')

    return parser

def validate(gt, disp):
    mask = (gt > 0)
    epe = np.abs(disp[mask] - gt[mask])
    error_1px = (epe > 1.0).sum() / np.count_nonzero(epe) * 100
    error_2px = (epe > 2.0).sum() / np.count_nonzero(epe) * 100
    error_3px = (epe > 3.0).sum() / np.count_nonzero(epe) * 100
    error_4px = (epe > 4.0).sum() / np.count_nonzero(epe) * 100
    error_5px = (epe > 5.0).sum() / np.count_nonzero(epe) * 100
    return {'epe': epe.mean(),
            'total_px': np.count_nonzero(epe),
            'error_1px': error_1px,
            'error_2px': error_2px,
            'error_3px': error_3px,
            'error_4px': error_4px,
            'error_5px': error_5px}

def visual_disparity(disp, img_name):
    disp_vis = (disp - disp.min()) / (disp.max() - disp.min()) * 255.0
    prob_histogram(disp_vis, img_name) # 显示概率直方图
    np.savetxt('./stereoflow/demo/' + img_name + '_pred_norm.txt', disp_vis, fmt='%.6f')
    disp_vis = cv2.applyColorMap(cv2.convertScaleAbs(disp_vis.astype("uint8"), alpha=-2), cv2.COLORMAP_JET) # 伪彩色映射
    return disp_vis

def prob_histogram(disp_vis, img_name):
    data = disp_vis.astype(np.uint8)
    flat_data = data.flatten()
    unique, counts = np.unique(flat_data, return_counts=True)
    probabilities = counts / len(flat_data)

    plt.bar(unique, probabilities)
    plt.savefig('./stereoflow/demo/' + img_name + '_pred_norm_histogram.png')
    # plt.show()

def main(args):
    dataset = args.dataset # kitti, usvinland
    img_name = args.imgname # 000000_10 for kitti, H05_1_0000000600 for usvinland

    gt = None
    if dataset == 'kitti':
        left_path = 'G:/KITTI/KITTI2015/data_scene_flow/training/image_2/' + img_name + '.png'
        if system == 'Ubuntu': left_path = left_path.replace('G:/', '/media/ubuntu/e/zhouyiqing/')
        image1 = np.asarray(Image.open(left_path))
        image2 = np.asarray(Image.open(left_path.replace('image_2', 'image_3')))
        gt = np.asarray(Image.open(left_path.replace('image_2', 'disp_occ_0'))) / 256.0
    elif 'usvinland' in dataset:
        left_path = 'G:/USVInland/Stereo Matching/Low_Res_640_320/Left_Img_Rectified/' + img_name + '.jpg'
        if system == 'Ubuntu': left_path = left_path.replace('G:/', '/media/ubuntu/e/zhouyiqing/')
        if dataset == 'usvinland_seg': left_path = left_path.replace('Low_Res_640_320', 'Segmentation').replace('Left_Img_Rectified', 'Left_Img_Rectified_Seg_rm')
        image1 = np.asarray(Image.open(left_path))
        image2 = np.asarray(Image.open(left_path.replace('Left', 'Right')))
        gt = np.asarray(Image.open(left_path.replace('Left_Img_Rectified', 'Disp_Map').replace('jpg', 'png'))) / 255.0 * 50.0

    assert os.path.exists(args.loadmodel), "Pth does not exist"
    model, _, cropsize, with_conf, task, tile_conf_mode = _load_model_and_criterion(args.loadmodel, None, device)

    im1 = img_to_tensor(image1).to(device).unsqueeze(0)
    im2 = img_to_tensor(image2).to(device).unsqueeze(0)
    with torch.inference_mode():
        pred, _, _ = tiled_pred(model, None, im1, im2, None, conf_mode=tile_conf_mode, overlap=tile_overlap, crop=cropsize, with_conf=with_conf, return_time=False)
    pred = pred.squeeze(0).squeeze(0).cpu().numpy()

    print('dataset:', dataset)
    print('img_name:', left_path.split('/')[-1])
    print('pred_shape:', pred.shape)
    if gt is not None:
        print('gt_shape:  ', gt.shape)
        error = validate(gt, pred)
        print('epe:', error['epe'])
        print('total_px:', error['total_px'])
        print('error_1px: {:6.4f}%'.format(error['error_1px']))
        print('error_2px: {:6.4f}%'.format(error['error_2px']))
        print('error_3px: {:6.4f}%'.format(error['error_3px']))
        print('error_4px: {:6.4f}%'.format(error['error_4px']))
        print('error_5px: {:6.4f}%'.format(error['error_5px']))
        np.savetxt('./stereoflow/demo/' + img_name + '_gt.txt', gt, fmt='%.6f')
        np.savetxt('./stereoflow/demo/' + img_name + '_pred.txt', pred, fmt='%.6f')

    pred_pseudo = visual_disparity(pred, img_name)
    cv2.imwrite('./stereoflow/demo/' + img_name + '_pred_pseudo.png', pred_pseudo)
    # cv2.imshow('stereo_demo', pred_pseudo)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

if __name__ == '__main__':
    args = get_args_parser()
    args = args.parse_args()
    main(args)
